"""
-------------------------------------------------
   File Name:       data_loader.py
   Project Name:    beamform_AYnet
   Author :         Chunshan YANG
   Date:            2025/2/2
   Device:          GTX2070
-------------------------------------------------
   Change Activity:
                   2025/2/2:
-------------------------------------------------
"""
import os
import os.path

import numpy as np
import torch
import torchvision.transforms as transforms
import torch.utils.data as data
import torch.nn.functional as F
from torch.utils.data import DataLoader
import scipy.io as scio
import skimage.io
import scipy.ndimage


def np_range_norm(image, maxminnormal=True, range1=True):

    if image.ndim == 2 or (image.ndim == 3 and image.shape[0] == 1):
        if maxminnormal:
            _min = image.min()
            _range = image.max() - image.min()
            narmal_image = (image - _min) / _range
            if range1:
               narmal_image = (narmal_image - 0.5) * 2
        else:
            _mean = image.mean()
            _std = image.std()
            narmal_image = (image - _mean) / _std

    return narmal_image

def pad_to_shape(array, target_shape):
    # 计算需要补零的行数和列数
    pad_height = target_shape[0] - array.shape[0]
    pad_width = target_shape[1] - array.shape[1]
        
    # 确保补零的行数和列数非负
    if pad_height < 0 or pad_width < 0:
        raise ValueError("Target shape must be larger than the current shape")
        
    # 使用 np.pad 进行补零
    padded_array = np.pad(array, ((0, pad_height), (0, pad_width)), mode = 'constant', constant_values = 0)
        
    return padded_array
    
def interpolate_to_shape(array, target_shape):
    # 计算每个维度的缩放因子
    zoom_factors = [target_shape[0] / array.shape[0], target_shape[1] / array.shape[1]]
        
    # 使用 scipy.ndimage.zoom 进行插值
    interpolated_array = scipy.ndimage.zoom(array, zoom_factors, order = 3)  # order=3 表示三次函数插值
        
    return interpolated_array
    
### 光声或超声数据集
class ReconDataset(data.Dataset):
    __inputdata = []   # RF data
    __inputimg = []    # beamform image
    __outputdata = []  # ground truth image

    def __init__(self, root, train=True, das=True, transform=None, modality='US'):
        self.__inputdata = []
        self.__outputdata = []
        self.__inputimg = []

        self.root = os.path.expanduser(root)
        self.transform = transform
        self.train = train
        self.modality = modality
        if train:
            folder = root + "Train/"
        else:
            folder = root + "Test/"
            
        for file in os.listdir(folder):
            print(file)

            # ## 如果文件不是mat文件，先找到对应的mat文件把acoustic_impedance或optical_absorption取出来用于归一化
            # if not file.endswith('_gt.png'):
            #     if file.endswith('_ac.mat'):
            #     elif file.endswith('_pa.mat'):
            #     else:
            #         print("input modality is not supported!")

            if file.endswith('.mat'):
                rf_data = scio.loadmat(folder + file)
                if modality == 'US':
                    sensor_data = rf_data['rf_data']
                    #self.acoustic_impedance = rf_data['acoustic_impedance']
                elif modality == 'PA':
                    sensor_data = rf_data['rf_data']
                    #self.optical_absorption = rf_data['optical_absorption']
                else:
                    print("input modality is not supported!")

                ## 将sensor_data_pa扩为128通道，2560个点
                sensor_data = pad_to_shape(sensor_data, (2560, 64))
                sensor_data = interpolate_to_shape(sensor_data, (2560, 128))
                sensor_data = np_range_norm(sensor_data, maxminnormal=True, range1=False)  ## 这种归一化方式是否把不同组织接收的信号强度不同都给消除了？
                sensor_data = sensor_data[np.newaxis, :, :]

                self.__inputdata.append(sensor_data)

            elif file.endswith('_gt.png'):
                gt_image = skimage.io.imread(folder + file)

                ## 将gt_image缩放为128 * 128
                gt_image = interpolate_to_shape(gt_image, (256, 128))
                gt_image = np_range_norm(gt_image, maxminnormal=True, range1=False)  ## 这儿问题不大，保存图像的过程中已经有归一化
                gt_image = gt_image[np.newaxis, :, :]
                
                self.__outputdata.append(gt_image)
                # # 显示图像
                # plt.figure(figsize=(10, 6))
                # plt.imshow(image, cmap='gray')  # 如果图像是灰度图像，使用 cmap='gray'
                # plt.title('Reconstructed US Image')
                # plt.axis('off')  # 关闭坐标轴
                # plt.show()

            elif das and file.endswith('_das.png'):
                das_image = skimage.io.imread(folder + file)

                ## 将gt_image缩放为256 * 128
                das_image = interpolate_to_shape(das_image, (256, 128))
                das_image = np_range_norm(das_image, maxminnormal=True, range1=False)   ## 这儿问题不大，保存图像的过程中已经有归一化
                das_image = das_image[np.newaxis, :, :]

                self.__inputimg.append(das_image)

            else:
                print("not support!")


            # matdata = scio.loadmat(folder + file)
            # self.__inputdata.append(np.transpose(matdata['sensor_data'])[np.newaxis,:,:])
            # self.__outputdata.append(matdata['p0'][np.newaxis,:,:])
            # if das:
            #     p0=np.delete(matdata['p0_das'],255,1)
            #     self.__inputimg.append(p0[np.newaxis,:,:])
            # else:
            #     self.__inputimg.append(matdata['p0_tr'][np.newaxis,:,:])

    def __getitem__(self, index):
        rawdata =  self.__inputdata[index] #.reshape((1,1,2560,120))
        #rawdata = (rawdata-(np.min(np.min(rawdata,axis=2)))/((np.max(np.max(rawdata,axis=2)))-(np.min(np.min(rawdata,axis=2))))
        #rawdata = rawdata -0.5
        #rawdata = np_range_norm(rawdata,maxminnormal=True)
        reconstruction = self.__outputdata[index] #.reshape((1,1,2560,120))
        #reconstruction = np_range_norm(reconstruction,maxminnormal=True)
        beamform = self.__inputimg[index]

        rawdata = torch.Tensor(rawdata)
        reconstructions = torch.Tensor(reconstruction)
        beamform = torch.Tensor(beamform)

        return rawdata, reconstructions, beamform

    def __len__(self):
        return len(self.__inputdata)
    
    def get_modality(self):
        return self.modality
    
    def get_acoustic_impedance(self):
        return self.acoustic_impedance
    
    def get_optical_absorption(self):
        return self.optical_absorption



if __name__ == "__main__":
    dataset_pathr = 'E:/YCS_Temp/project/thesis_dataset/dataset/PA_DATA/'

    mydataset = ReconDataset(dataset_pathr, train = False, das = True, modality = 'PA')
    #print(mydataset.__getitem__(3))
    train_loader = DataLoader(mydataset, batch_size=1, shuffle = True)
    batch_idx, (rawdata, reimage, bfim) = list(enumerate(train_loader))[0]
    print(rawdata.size())
    print(rawdata.max())
    print(rawdata.min())
    print(mydataset.__len__())






